import json
import os
import struct
import argparse
from typing import BinaryIO, Tuple, List


class DayData:
  """
  日线数据类，用于存储某只股票一天的交易数据
  """
  def __init__(self, index: int, data: bytes, address: int, days_data: List['DayData']):
    self.index = index                  # 数据索引
    self.days_data = days_data          # 所有日线数据列表
    self.address = address              # 数据在文件中的地址
    self.data = data                    # 原始字节数据
    self.date = struct.unpack('<i', data[:4])[0]           # 日期（Unix时间戳格式）
    self.open = round(struct.unpack('<f', data[8:12])[0], 4)   # 开盘价
    self.close = round(struct.unpack('<f', data[12:16])[0], 4) # 收盘价
    self.high = round(struct.unpack('<f', data[16:20])[0], 4)  # 最高价
    self.low = round(struct.unpack('<f', data[20:24])[0], 4)   # 最低价
    self.volume = struct.unpack('<i', data[24:28])[0]         # 成交量
    self.unknown_i = struct.unpack('<i', data[36:40])[0]      # 未知整数字段
    self.unknown_f = struct.unpack('<f', data[36:40])[0]      # 未知浮点数字段

  def __repr__(self):
    return f'<DataChunk index={self.index} {hex(self.address)} date={self.date} o={self.open}, c={self.close} h={self.high} l={self.low} v={self.volume}>'

  def to_dict(self):
    """
    将日线数据转换为字典格式，便于JSON序列化
    """
    return {
      'date': str(self.date), 
      'open': self.open, 
      'close': self.close, 
      'high': self.high, 
      'low': self.low,
      'volume': self.volume,
      # 上一个交易日收盘价，如果是第一条数据则使用当前收盘价
      'last_close': self.days_data[self.index - 1].close if self.index > 0 else self.close
    }


class DataChunk:
  """
  数据块类，包含400条日线数据记录
  """
  def __init__(self, index: int, data: bytes, address: int, chunk_index_with_code: Tuple[int, 'Code']):
    self.index = index                            # 数据块索引
    self.address = address                        # 数据块在文件中的地址
    self.code_chunk_id, self.code = chunk_index_with_code  # 代码块ID和对应的股票代码
    self.days_data = []                           # 日线数据列表
    
    # 计算当前数据块包含的日线数据条数
    # 如果不是最后一个数据块，则包含400条记录；否则计算余数
    day_count = 400 if self.code_chunk_id < len(self.code.chunk_indexs) - 1 else (self.code.data_count % 400)
    
    # 解析日线数据
    for i in range(day_count):
      day_data = DayData(i, data[i*40:(i+1)*40], address + i*40, self.days_data)
      self.days_data.append(day_data)
    
    # 将解析出的日线数据添加到对应的股票代码中
    self.code.days_data.extend(self.days_data)

  def __repr__(self):
    return f'<DataChunk index={self.index} {hex(self.address)} code={self.code.code} days_data:frist={self.days_data[0]} last={self.days_data[-1]}>'


class DataChunks:
  """
  所有数据块的集合类，负责解析整个数据文件中的所有日线数据块
  """
  def __init__(self, f: BinaryIO, head: 'Head', codes: 'Codes'):
    self.data_chunks = []     # 数据块列表
    self.codes = codes        # 股票代码信息
    
    # 定位到数据块起始位置
    # 偏移量 = 代码区域大小 + 头部大小 + 8字节偏移
    f.seek(head.codes_space_count * 0x204 + 0x28 + 8)
    
    # 依次读取并解析每个数据块
    for i in range(head.chunk_count):
      tell = f.tell()  # 当前文件指针位置
      # 每个数据块大小为40*400字节（每条日线数据40字节，每块400条）
      data_chunk = DataChunk(i, f.read(40 * 400), tell, self.codes.chunk_index_code[i])
      self.data_chunks.append(data_chunk)
    print('data_chunk=', self.data_chunks[-1])


class Code:
  """
  股票代码类，存储股票代码及其相关数据信息
  """
  def __init__(self, index: int, data: bytes, address: int):
    self.index = index                    # 股票索引
    self.address = address                # 数据在文件中的地址
    self.code = data[8:14].decode('ascii')  # 股票代码（如：000001）
    self.data_count = struct.unpack('<i', data[32:36])[0]  # 该股票的日线数据总数
    self.chunk_indexs = []                # 数据块索引列表
    self.days_data = []  # type: List[DayData]  # 日线数据列表
    
    # 解析数据块索引
    start_pos = 36
    # 当未遇到结束标记时，继续读取数据块索引
    while data[start_pos: start_pos+4] != b'\xff\xff\xff\xff':
      chunk_index = struct.unpack('<i', data[start_pos:start_pos+4])[0]
      self.chunk_indexs.append(chunk_index)
      start_pos += 4

  def __repr__(self):
    return f'<Code index={self.index} {hex(self.address)} code={self.code} data_count={self.data_count} chunk_indexs={self.chunk_indexs}>'

  def to_dict(self):
    return [day_data.to_dict() for day_data in self.days_data]

  def save_json(self, save_path):
    """
    将股票数据保存为JSON文件
    :param save_path: 保存路径
    """
    with open(save_path, 'w', encoding='utf8') as f:
      f.write(json.dumps(self.to_dict(), indent=2, ensure_ascii=False))


class Codes:
  """
  股票代码集合类，管理所有股票代码信息
  """
  def __init__(self, f: BinaryIO, head: 'Head'):
    self.codes = []                # 股票代码列表
    self.chunk_index_code = {}     # 数据块索引与股票代码映射字典
    
    # 依次读取并解析每只股票的代码信息
    for i in range(head.stocks_count):  # 股票总数
      tell = f.tell()
      # 每个股票代码信息占0x204字节
      code = Code(i, f.read(0x204), tell)
      self.codes.append(code)
      # print('code=', code)
      
      # 建立数据块索引与股票代码的映射关系
      for j, index in enumerate(code.chunk_indexs):
        self.chunk_index_code[index] = (j, code)


class Head:
  """
  文件头信息类，存储数据文件的基本信息
  """
  def __init__(self, data: bytes):
    print(f'data={data.hex(" ")}')
    self.version = struct.unpack('<ii', data[:8])              # 版本信息
    self.stocks_count = struct.unpack('<i', data[8:12])[0]     # 股票数量
    self.chunk_count = struct.unpack('<i', data[12:16])[0]     # 数据块数量
    self.stock_chunk_size = struct.unpack('<i', data[0x10:0x14])[0]  # 股票数据块大小
    self.codes_space_count = struct.unpack('<i', data[0x14:0x18])[0] # 代码区域占用空间数量
    self.unknown_chunk_18 = struct.unpack('<i', data[0x18:0x1c])[0]  # 未知字段
    self.unknown_chunk_1c = struct.unpack('<i', data[0x1c:0x20])[0]  # 未知字段

  def __repr__(self):
    return f'<Head version={self.version} stocks_count={self.stocks_count} chunk_count={self.chunk_count} stock_chunk_size={self.stock_chunk_size} codes_space_count={self.codes_space_count} unknown:18={self.unknown_chunk_18},1c={self.unknown_chunk_1c}>'


class Reader:
  """
  数据文件读取器主类，负责解析整个数据文件
  """
  def __init__(self, file_path):
    self.head = None  # type: Head | None       # 文件头信息
    self.codes = None  # type: Codes | None     # 股票代码信息
    self.data_chunks = None  # type: DataChunks | None  # 数据块信息
    
    # 以二进制模式打开文件并开始解析
    with open(file_path, 'rb') as f:
      self.f = f
      self._analyze()

  def _analyze(self):
    """
    分析整个数据文件，依次解析头部、代码和数据块信息
    """ 
    self.head = self._analyze_head()           # 解析文件头
    print('self.head=', self.head)
    self.codes = self._analyze_codes()         # 解析股票代码
    print('self.codes=', self.codes)
    self.data_chunks = self._analyze_data_chunks()  # 解析数据块
    print('self.data_chunks=', self.data_chunks)

  def _analyze_head(self, length=0x28):
    return Head(self.f.read(length))

  def _analyze_codes(self):
    return Codes(self.f, self.head)

  def _analyze_data_chunks(self):
    return DataChunks(self.f, self.head, self.codes)

  def save(self, save_path, white_codes=None):
    """
    保存解析后的数据为JSON文件
    :param save_path: 保存目录路径
    :param white_codes: 白名单股票代码列表，如果指定则只保存白名单内的股票数据
    """
    os.makedirs(save_path, exist_ok=True)  # 创建保存目录（如果不存在）
    
    # 遍历所有股票代码并保存数据
    for code in self.codes.codes:
      # 如果指定了白名单且当前股票不在白名单内，则跳过
      if white_codes is not None and code.code not in white_codes:
        continue
      code.save_json(f'{save_path}/{code.code}.json')  # 保存为JSON文件


def main(file_path, save_path, save_white_path=None):
  """
  主函数，解析本地数据文件并保存为JSON格式
  :param file_path: 数据文件路径
  :param save_path: JSON数据保存路径
  :param save_white_path: 白名单文件路径（可选）
  """
  white_codes = None
  # 如果指定了白名单文件，则读取白名单股票代码
  if save_white_path is not None:
    with open(save_white_path, encoding='utf8') as f:
      white_codes = f.read().split('\n')
  
  # 创建读取器并解析数据文件
  reader = Reader(file_path)
  # 保存解析后的数据
  reader.save(save_path, white_codes)


if __name__ == '__main__':
  """
  程序入口点，支持命令行参数
  """
  parser = argparse.ArgumentParser(description='日线V43解析器')
  parser.add_argument('input_file', help='日线V43数据文件路径')
  parser.add_argument('output_dir', help='输出目录路径')
  parser.add_argument('--whitelist', '-w', help='白名单文件路径（可选）')
  
  args = parser.parse_args()
  
  main(args.input_file, args.output_dir, args.whitelist)
